package com.aikg.kgrag_java.service.impl;

import cn.hutool.core.util.StrUtil;
import co.elastic.clients.elasticsearch.ElasticsearchClient;
import com.aikg.kgrag_java.common.Result;
import co.elastic.clients.elasticsearch._types.mapping.TypeMapping;
import co.elastic.clients.elasticsearch._types.query_dsl.Query;
import co.elastic.clients.elasticsearch._types.query_dsl.RangeQuery;
import co.elastic.clients.elasticsearch.core.*;

import co.elastic.clients.elasticsearch.core.search.Hit;
import co.elastic.clients.elasticsearch.indices.*;
import co.elastic.clients.json.JsonData;
import co.elastic.clients.util.ObjectBuilder;
import com.aikg.kgrag_java.pojo.entity.IndexDocument;
import com.aikg.kgrag_java.service.IndexService;
import com.alibaba.fastjson2.JSONObject;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.elasticsearch.client.*;
import org.springframework.stereotype.Service;

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.function.Function;

/**
 * 实现 IndexService 接口，提供 Elasticsearch 索引的增删改查操作。
 */
@Service
@Slf4j
public class IndexServiceImpl implements IndexService {
    @Resource
    private ElasticsearchClient elasticsearchClient;

    @Resource
    private RestClient restClient;

    /**
     * 创建一个新的索引库
     *
     * @param name 索引的名称
     * @return
     * @throws IOException 如果创建索引过程中出现 IO 异常
     */
    @Override
    public Result addIndex(String name) throws IOException {
        if (StrUtil.isBlank(name)) {
            return Result.buildErrorResult("索引名称不能为空");
        }

        // 判断索引是否存在
        if (indexExists(name)) {
            return Result.buildErrorResult("索引已存在");
        }

        try{
            // 使用 Elasticsearch 客户端创建索引，调用 `indices().create` 方法创建指定名称的索引
            CreateIndexResponse createIndexResponse = elasticsearchClient.indices().create(c -> c.index(name));
            if( createIndexResponse.acknowledged()){
                return Result.buildResult(200, "创建索引成功");
            }else {
                return Result.buildErrorResult("创建索引失败");
            }
        }catch (IOException e){
            return Result.buildErrorResult("创建索引失败：" + e.getMessage());
        }
    }

    /**
     * 判断指定的索引是否已存在。
     *
     * @param name 索引的名称
     * @return 如果索引存在，返回 true；否则返回 false
     * @throws IOException 如果在检查索引存在性过程中出现 IO 异常
     */
    @Override
    public boolean indexExists(String name) throws IOException {
        // 使用 Elasticsearch 客户端检查索引是否存在，调用 `indices().exists` 方法
        return elasticsearchClient.indices().exists(b -> b.index(name)).value();
    }

    /**
     * 删除指定的索引。
     *
     * @param name 索引的名称
     * @return
     * @throws IOException 如果删除索引过程中出现 IO 异常
     */
    @Override
    public Result delIndex(String name) throws IOException {
        if (StrUtil.isBlank(name)) {
            log.error("索引名称不能为空");
            return Result.buildErrorResult("索引名称不能为空");
        }
        // 判断索引是否存在
//        if (!indexExists(name)) {
//            log.error("索引不存在");
//            return Result.buildErrorResult("索引不存在");
//        }

        try {
            // 使用 Elasticsearch 客户端删除指定名称的索引，调用 `indices().delete` 方法
            DeleteIndexResponse delete = elasticsearchClient.indices().delete(c -> c.index(name));

            // 记录日志并返回结果
            if (delete.acknowledged()) {
                return Result.buildResult(200,"索引删除成功");
            } else {
                return Result.buildErrorResult("索引删除失败");
            }
        } catch (IOException e) {
            return Result.buildErrorResult("删除索引过程中发生异常: " + e.getMessage());
        }
    }

    /**
     * 创建带有指定设置和映射的索引。
     *
     * @param name      索引的名称
     * @param settingFn 提供索引设置的函数，返回 IndexSettings 对象的构建器
     * @param mappingFn 提供索引映射的函数，返回 TypeMapping 对象的构建器
     * @return
     * @throws IOException 如果创建索引过程中出现 IO 异常
     */
    @Override
    public boolean create(String name,
                          Function<IndexSettings.Builder, ObjectBuilder<IndexSettings>> settingFn,
                          Function<TypeMapping.Builder, ObjectBuilder<TypeMapping>> mappingFn) throws IOException {

        if(indexExists(name)){
            log.error("索引已存在");
            return false;
        }
        // 使用 CreateIndexRequest 来构建索引请求
        CreateIndexRequest.Builder builder = new CreateIndexRequest.Builder();
        builder.index(name);

        // 设置索引的设置和映射
        builder.settings(settingFn);
        builder.mappings(mappingFn);

        // 构建请求对象
        CreateIndexRequest request = builder.build();

        // 执行索引创建请求
        CreateIndexResponse response = elasticsearchClient.indices().create(request);

        return response.acknowledged();
    }

    /**
     * 查询索引详情
     *
     * @param name
     * @return
     * @throws IOException
     */
    @Override
    public GetIndexResponse getIndexDetail(String name) throws IOException {
        //使用 * 或者 _all都可以
        GetIndexResponse response = elasticsearchClient.indices().get(builder -> builder.index(name));
        log.info("getIndexDetail方法，response.result()={}", response.result().toString());
        return response;
    }

    @Override
    /**
     * 向指定索引插入数据
     *
     * @param indexName 索引名称
     * @param id        文档的唯一 ID
     * @param document  要插入的数据对象
     * @return 插入操作是否成功
     * @throws IOException 如果插入数据过程中出现 IO 异常
     */
    public boolean insertDocument(String indexName, String id, IndexDocument document) throws IOException {

        if(StrUtil.isBlank(indexName) || StrUtil.isBlank(id) || document == null){
            return false;
        }

        // 将 IndexDocument 对象转换为 Map 或 POJO，准备插入 Elasticsearch
        Map<String, Object> source = convertEntityToSource(document);

        // 创建一个 IndexRequest 来执行插入操作
        IndexRequest indexRequest = new IndexRequest.Builder()
                .index(indexName)  // 指定索引名称
//                .id(id)  // 指定文档 ID
                .document(source)  // 指定文档内容（从 IndexDocument 转换的 Map）
                .build();

        // 使用 Elasticsearch 客户端执行插入操作
        IndexResponse response = elasticsearchClient.index(indexRequest);

        // 根据响应结果判断插入是否成功
        co.elastic.clients.elasticsearch._types.Result result = response.result();

        // 根据响应结果判断插入是否成功
        if (result == co.elastic.clients.elasticsearch._types.Result.Created || result == co.elastic.clients.elasticsearch._types.Result.Updated) {
            log.info("成功插入文档，索引: {}, ID: {}", indexName, id);
            return true;
        } else {
            log.error("插入文档失败，索引: {}, ID: {}", indexName, id);
            return false;
        }
    }


    /**
     * 将 IndexDocument 转换为 Elasticsearch 插入所需的 Map 对象
     *
     * @param document IndexDocument 文档对象
     * @return 转换后的 Map
     */
    private Map<String, Object> convertEntityToSource(IndexDocument document) {
        // 将 IndexDocument 的字段转换为 Map 格式
        return Map.of(
                "project_id", document.getProjectId(),
                "domain_name", document.getDomainName(),
                "knowledgedb", document.getKnowledgedb(),
                "source", document.getSource(),
                "page_content", document.getPageContent(),
                "vector", document.getVector(), // 向量字段
                "chunk_number", document.getChunkNumber()
        );
    }

    /**
     * 使用余弦相似度进行 k-NN 向量检索
     *通过余弦相似度计算查询向量与文档向量之间的相似度，返回最相似的 k 个文档。
     * @param queryVector 查询向量
     * @param k           表示返回前 k 个最相似的文档
     * @param indexName   在同一个索引库里检索，即在同一个领域
     * @return
     * @throws IOException
     */
    public List<IndexDocument> searchByCosineSimilarity(List<Double> queryVector, int k, String indexName, String field) throws IOException {
        if(queryVector == null || queryVector.size() == 0 || StrUtil.isBlank(indexName) || StrUtil.isBlank(field) || k <= 0){
            return null;
        }

        // 创建POST请求
        Request postRequest = new Request("POST", indexName+"/_search");

        String jsonBody = "{\n" +
                "  \"size\": 3,\n" +
                "  \"query\": {\n" +
                "    \"script_score\": {\n" +
                "      \"query\": {\n" +
                "        \"match_all\": {}\n" +
                "      },\n" +
                "      \"script\": {\n" +
                "        \"source\": \"cosineSimilarity(params.query_vector, 'vector') + 1.0\",\n" +
                "        \"params\": {\n" +
                "          \"query_vector\": " + queryVector + "\n" +
                "        }\n" +
                "      }\n" +
                "    }\n" +
                "  },\n" +
                "  \"sort\": [\n" +
                "    { \"_score\": { \"order\": \"desc\" } }\n" +
                "  ]\n" +
                "}";
        Response response;
        // 设置请求体
        postRequest.setJsonEntity(jsonBody);
        // 防止查询失败返回400~500错误报错
        try {
            response = restClient.performRequest(postRequest);
        } catch (ResponseException e) {
            return new ArrayList<>();
        }

//        System.out.println("status:" + response.getStatusLine().getStatusCode());

        // 解析响应
        BufferedReader bufferedReader = new BufferedReader(new InputStreamReader(response.getEntity().getContent()));
        StringBuilder responseBody = new StringBuilder();

        String str = null;
        while ((str = bufferedReader.readLine()) != null) {
            responseBody.append(str);
        }
        bufferedReader.close();

        // 获取 JSON 响应字符串
        String jsonResponse = responseBody.toString();

        // 使用 Jackson 将响应映射为 JsonNode
        ObjectMapper objectMapper = new ObjectMapper();
        JsonNode responseNode = objectMapper.readTree(jsonResponse);

        // 获取 'hits' 节点，它是一个数组
        JsonNode hitsNode = responseNode.path("hits").path("hits");

        List<IndexDocument> documents = new ArrayList<>();

        // 遍历 'hits' 数组
        if (hitsNode.isArray()) {
            for (JsonNode hitNode : hitsNode) {
                JsonNode sourceNode = hitNode.path("_source");

                // 将返回的数据映射到 IndexDocument 对象
                IndexDocument document = new IndexDocument();
                document.setProjectId(sourceNode.path("project_id").asText());
                document.setDomainName(sourceNode.path("domain_name").asText());
                document.setKnowledgedb(sourceNode.path("knowledgedb").asText());
                document.setSource(sourceNode.path("source").asText());
                document.setPageContent(sourceNode.path("page_content").asText());

                // 对向量进行解析，如果有的话
                JsonNode vectorNode = sourceNode.path("vector");
                if (vectorNode.isArray()) {
                    List<Double> vector = new ArrayList<>();
                    for (JsonNode vectorElement : vectorNode) {
                        vector.add(vectorElement.asDouble());
                    }
                    document.setVector(vector);
                }

                document.setChunkNumber(sourceNode.path("chunk_number").asText());

                documents.add(document);
            }
        }

        return documents;
    }

    /**
     * 删除有条件删除文档数据
     * @param indexName
     * @param fieldName
     * @param value
     * @return
     * @param <T>
     * @throws IOException
     */
    public <T> boolean deleteByQuery(String indexName, String fieldName,T value) {

        if (indexName == null || StrUtil.isBlank(indexName)) {
            throw new IllegalArgumentException("索引不能为空");
        }

        if (fieldName == null || StrUtil.isBlank(fieldName)) {
            throw new IllegalArgumentException("fieldName字段为空");
        }

        // 根据value的类型动态构建查询条件
        Query query = buildQuery(fieldName, value);

        // 构建DeleteByQueryRequest,删除指定索引库里面的文档数据
        DeleteByQueryRequest request = new DeleteByQueryRequest.Builder()
                .index(indexName)
                .query(query)
                .build();


//         执行删除操作
        try {
            DeleteByQueryResponse response = elasticsearchClient.deleteByQuery(request);
            if (response.deleted() > 0) {
                return true;
            }
        } catch (Exception e) {
            log.error(e.getMessage());
            throw new RuntimeException("es文档删除失败");
        }
        return false;
    }

    /**
     * 构建查询条件，支持不同类型的值
     *
     * @param fieldName 字段名称
     * @param value 查询值
     * @param <T> 查询值的类型
     * @return Query
     */
    private <T> Query buildQuery(String fieldName, T value) {
        // 判断value的类型，并构建不同的查询条件
        if (value instanceof String) {
            return Query.of(q -> q
                    .term(t -> t
                            .field(fieldName)
                            .value((String) value)  // 将 value 强制转换为 String
                    )
            );
        } else if (value instanceof Long) {
            return Query.of(q -> q
                    .term(t -> t
                            .field(fieldName)
                            .value((Long) value)  // 将 value 强制转换为 Long
                    )
            );
        } else if (value instanceof Double) {
            return Query.of(q -> q
                    .term(t -> t
                            .field(fieldName)
                            .value((Double) value)  // 将 value 强制转换为 Double
                    )
            );
        } else if (value instanceof Boolean) {
            return Query.of(q -> q
                    .term(t -> t
                            .field(fieldName)
                            .value((Boolean) value)  // 将 value 强制转换为 Boolean
                    )
            );
        } else {
            throw new IllegalArgumentException("Unsupported value type: " + value.getClass().getName());
        }
    }

}
